from klineapp.models import SyncSessionFactory
from common.log import LOG, LogLevel
import akshare as ak
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
import asyncio


def make_main_table(futures):
    """创建一张主表，上面存储所有合约的代码"""
    session = SyncSessionFactory()  # 获取同步会话
    try:
        with session.begin():
            # 创建主表
            futures.to_sql(
                name="main_contracts",
                con=session.bind,  # 这里使用线程自己的同步session
                if_exists="replace",  # 如果表存在就替换
                index=False,
            )
            LOG(LogLevel.INFO) << "主表 'main_contracts' 已创建，包含所有主力合约代码！" << "\n"
    except Exception as e:
        LOG(LogLevel.ERROR) << f"创建主表失败: {str(e)}" << "\n"
        raise
    finally:
        session.close()


def process_one_symbol(row):
    """处理单个合约的函数"""
    session = SyncSessionFactory()  # 获取同步会话
    try:
        with session.begin():
            periods = ["1", "5", "30", "120"] # 定义需要获取的周期，单位为min
            for period in periods:
                # 获取数据...（只拿最近100根k线数据）
                kline_data = ak.futures_zh_minute_sina(symbol=row.symbol, period=period).tail(100)
                # 存储数据...
                kline_data.to_sql(
                    name=f"{row.symbol.lower()}_{period}min",
                    con=session.bind,  # 这里使用线程自己的同步session
                    if_exists="fail",  # 如果表存在就抛出异常
                    index=False,
                )
            return (row.symbol, True)
    except Exception as e:
        tqdm.write(f"处理合约 {row.symbol} 时出错: {str(e)}")
        return (row.symbol, False)
    finally:
        session.close()  # 确保关闭session


async def process_symbols_concurrently(futures):
    """并发处理所有合约"""
    loop = asyncio.get_running_loop()
    success_count = 0  # 成功计数器
    
    with tqdm(total=len(futures), desc="数据入库进度") as pbar:
        # 根据自己电脑的CPU情况开启线程池（CPU核心数*2）
        with ThreadPoolExecutor(max_workers=4) as pool:
            # 提交所有任务到线程池
            tasks = [loop.run_in_executor(pool, process_one_symbol, row) 
                    for row in futures.itertuples()]
            # 处理结果
            for task in asyncio.as_completed(tasks):
                symbol, result = await task
                if result is True:
                    success_count += 1
                    pbar.update(1)
                    pbar.set_postfix_str(f"完成: {symbol}")
                else:
                    pbar.set_postfix_str(f"失败: {symbol}")
                pbar.refresh()

    # 检查是否所有任务都成功完成
    if success_count != len(futures):
        failed_count = len(futures) - success_count
        raise RuntimeError(f"部分数据未加载成功，失败数量: {failed_count}/{len(futures)}")